from langchain.agents import load_tools, initialize_agent, AgentType, AgentExecutor
from langchain.agents.format_scratchpad import format_to_openai_function_messages
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser
from langchain.chat_models import ChatOpenAI
from langchain.llms.openai import OpenAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.schema import AgentFinish, HumanMessage, AIMessage
from langchain.tools import tool
from langchain.tools.render import format_tool_to_openai_function


def normal_agent():
    llm = OpenAI(temperature=0)
    tools = load_tools(["serpapi", "llm-math"], llm)
    agent = initialize_agent(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True)
    result = agent.run("Who is Leo DiCaprio`s girlfriend? What isher current age raised to the 0.43 power?")
    print(result)


@tool
def get_word_length(word: str) -> int:
    """Returns the length of a word."""
    return len(word)


def get_word_length_agent():
    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, verbose=True)
    tools = [get_word_length]
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "You are very powerful assistant, but bad at calculating lengths of words.",
            ),
            ("user", "{input}"),
            MessagesPlaceholder(variable_name="agent_scratchpad"),
        ]
    )
    llm_with_tools = llm.bind(functions=[format_tool_to_openai_function(t) for t in tools])
    agent = (
            {
                "input": lambda x: x["input"],
                "agent_scratchpad": lambda x: format_to_openai_function_messages(
                    x["intermediate_steps"]
                ),
            }
            | prompt
            | llm_with_tools
            | OpenAIFunctionsAgentOutputParser()
    )
    return agent


def custom_agent():
    agent = get_word_length_agent()
    print(agent.invoke({"input": "how many letters in the word educa?", "intermediate_steps": []}))


def define_the_runtime():
    agent = get_word_length_agent()
    user_input = "how many letters in the word educa?"
    intermediate_steps = []
    while True:
        output = agent.invoke(
            {
                "input": user_input,
                "intermediate_steps": intermediate_steps,
            }
        )
        if isinstance(output, AgentFinish):
            final_result = output.return_values["output"]
            break
        else:
            print(f"TOOL NAME: {output.tool}")
            print(f"TOOL INPUT: {output.tool_input}")
            tool = {"get_word_length": get_word_length}[output.tool]
            observation = tool.run(output.tool_input)
            intermediate_steps.append((output, observation))
    print(final_result)


def add_memory_to_agent():
    tools = [get_word_length]
    agent = get_word_length_agent()
    chat_history = []
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    input1 = "how many letters in the word educa?"
    result = agent_executor.invoke({"input": input1, "chat_history": chat_history})
    chat_history.extend(
        [
            HumanMessage(content=input1),
            AIMessage(content=result["output"]),
        ]
    )
    agent_executor.invoke({"input": "is that a real word?", "chat_history": chat_history})


if __name__ == '__main__':
    add_memory_to_agent()

